"""
数据源整合模块

负责从各种数据源收集数据并整合为统一格式，包括:
- 历史价格数据
- 交易量数据
- 技术指标数据
- 新闻情感分析结果
"""

import os
import logging
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Union, Tuple, Any

# 配置日志
logger = logging.getLogger(__name__)

class PriceDataSource:
    """历史价格数据源"""
    
    def __init__(self, api_key: str):
        """
        初始化价格数据源
        
        Args:
            api_key: API访问密钥
        """
        self.api_key = api_key
        logger.info("初始化价格数据源")
        
    def get_daily_prices(self, stock_code: str, start_date: str, end_date: str) -> Dict[str, Dict[str, float]]:
        """
        获取股票的每日价格数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            字典格式的价格数据，键为日期，值为包含开盘价、收盘价、最高价、最低价的字典
        """
        logger.info(f"获取股票 {stock_code} 从 {start_date} 到 {end_date} 的每日价格数据")
        
        # TODO: 实现实际数据获取逻辑，这里使用模拟数据
        try:
            # 这里应该是实际从API获取数据的代码
            # 目前使用随机数据模拟
            date_range = pd.date_range(start=start_date, end=end_date)
            base_price = np.random.uniform(50, 200)
            
            price_data = {}
            for date in date_range:
                date_str = date.strftime("%Y-%m-%d")
                daily_volatility = np.random.uniform(0.98, 1.02)
                open_price = base_price * daily_volatility
                close_price = open_price * np.random.uniform(0.98, 1.02)
                high_price = max(open_price, close_price) * np.random.uniform(1.0, 1.01)
                low_price = min(open_price, close_price) * np.random.uniform(0.99, 1.0)
                
                price_data[date_str] = {
                    "open": round(open_price, 2),
                    "close": round(close_price, 2),
                    "high": round(high_price, 2),
                    "low": round(low_price, 2)
                }
                
                # 更新基准价格为当天收盘价
                base_price = close_price
            
            return price_data
        
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 价格数据时出错: {str(e)}")
            return {}


class VolumeDataSource:
    """交易量数据源"""
    
    def __init__(self, api_key: str):
        """
        初始化交易量数据源
        
        Args:
            api_key: API访问密钥
        """
        self.api_key = api_key
        logger.info("初始化交易量数据源")
        
    def get_daily_volumes(self, stock_code: str, start_date: str, end_date: str) -> Dict[str, int]:
        """
        获取股票的每日交易量数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            字典格式的交易量数据，键为日期，值为交易量
        """
        logger.info(f"获取股票 {stock_code} 从 {start_date} 到 {end_date} 的每日交易量数据")
        
        # TODO: 实现实际数据获取逻辑，这里使用模拟数据
        try:
            # 这里应该是实际从API获取数据的代码
            # 目前使用随机数据模拟
            date_range = pd.date_range(start=start_date, end=end_date)
            base_volume = np.random.randint(100000, 1000000)
            
            volume_data = {}
            for date in date_range:
                date_str = date.strftime("%Y-%m-%d")
                daily_volume = int(base_volume * np.random.uniform(0.8, 1.2))
                volume_data[date_str] = daily_volume
                
                # 随机更新基准交易量
                if np.random.random() < 0.2:  # 20%的概率
                    base_volume = int(base_volume * np.random.uniform(0.9, 1.1))
            
            return volume_data
        
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 交易量数据时出错: {str(e)}")
            return {}


class TechnicalIndicatorsCalculator:
    """技术指标计算器"""
    
    def __init__(self):
        """初始化技术指标计算器"""
        logger.info("初始化技术指标计算器")
    
    def calculate_ma(self, prices: List[float], window: int) -> List[float]:
        """
        计算移动平均线
        
        Args:
            prices: 价格序列
            window: 窗口大小
            
        Returns:
            移动平均线值列表
        """
        if len(prices) < window:
            return [np.nan] * len(prices)
        
        ma = []
        for i in range(len(prices)):
            if i < window - 1:
                ma.append(np.nan)
            else:
                ma.append(sum(prices[i-window+1:i+1]) / window)
        
        return ma
    
    def calculate_rsi(self, prices: List[float], window: int = 14) -> List[float]:
        """
        计算相对强弱指数 (RSI)
        
        Args:
            prices: 价格序列
            window: 窗口大小，默认为14
            
        Returns:
            RSI值列表
        """
        if len(prices) <= window:
            return [np.nan] * len(prices)
        
        # 计算价格变化
        deltas = np.diff(prices)
        seed = deltas[:window+1]
        up = seed[seed >= 0].sum() / window
        down = -seed[seed < 0].sum() / window
        
        if down == 0:
            rs = np.inf
        else:
            rs = up / down
        
        rsi = np.zeros_like(prices)
        rsi[window] = 100. - 100. / (1. + rs)
        
        # 计算剩余的RSI值
        for i in range(window+1, len(prices)):
            delta = deltas[i-1]
            
            if delta > 0:
                upval = delta
                downval = 0
            else:
                upval = 0
                downval = -delta
            
            up = (up * (window - 1) + upval) / window
            down = (down * (window - 1) + downval) / window
            
            if down == 0:
                rs = np.inf
            else:
                rs = up / down
            
            rsi[i] = 100. - 100. / (1. + rs)
        
        # 填充前window个值为NaN
        rsi[:window] = np.nan
        
        return rsi.tolist()
    
    def calculate_macd(self, prices: List[float], fast: int = 12, slow: int = 26, signal: int = 9) -> Tuple[List[float], List[float], List[float]]:
        """
        计算移动平均收敛/发散(MACD)
        
        Args:
            prices: 价格序列
            fast: 快速EMA窗口大小，默认为12
            slow: 慢速EMA窗口大小，默认为26
            signal: 信号线的窗口大小，默认为9
            
        Returns:
            包含MACD线、信号线和柱状图值的元组
        """
        # 将价格转换为numpy数组
        prices_array = np.array(prices)
        
        # 计算EMA
        ema_fast = np.zeros_like(prices_array)
        ema_slow = np.zeros_like(prices_array)
        
        # 初始化EMA
        ema_fast[slow-1] = np.mean(prices_array[:slow])
        ema_slow[slow-1] = ema_fast[slow-1]
        
        # 计算剩余的EMA值
        for i in range(slow, len(prices_array)):
            ema_fast[i] = (prices_array[i] - ema_fast[i-1]) * (2 / (fast + 1)) + ema_fast[i-1]
            ema_slow[i] = (prices_array[i] - ema_slow[i-1]) * (2 / (slow + 1)) + ema_slow[i-1]
        
        # 计算MACD线
        macd_line = ema_fast - ema_slow
        
        # 计算信号线
        signal_line = np.zeros_like(macd_line)
        signal_line[slow+signal-2] = np.mean(macd_line[slow-1:slow+signal-1])
        
        for i in range(slow+signal-1, len(macd_line)):
            signal_line[i] = (macd_line[i] - signal_line[i-1]) * (2 / (signal + 1)) + signal_line[i-1]
        
        # 计算柱状图
        histogram = macd_line - signal_line
        
        # 填充前slow+signal-2个值为NaN
        macd_line[:slow-1] = np.nan
        signal_line[:slow+signal-2] = np.nan
        histogram[:slow+signal-2] = np.nan
        
        return macd_line.tolist(), signal_line.tolist(), histogram.tolist()
    
    def calculate_bollinger_bands(self, prices: List[float], window: int = 20, num_std: float = 2.0) -> Tuple[List[float], List[float], List[float]]:
        """
        计算布林带
        
        Args:
            prices: 价格序列
            window: 窗口大小，默认为20
            num_std: 标准差的倍数，默认为2.0
            
        Returns:
            包含中轨、上轨和下轨的元组
        """
        if len(prices) < window:
            return ([np.nan] * len(prices), [np.nan] * len(prices), [np.nan] * len(prices))
        
        # 计算移动平均线(中轨)
        middle_band = self.calculate_ma(prices, window)
        
        # 计算标准差
        rolling_std = []
        for i in range(len(prices)):
            if i < window - 1:
                rolling_std.append(np.nan)
            else:
                std = np.std(prices[i-window+1:i+1])
                rolling_std.append(std)
        
        # 计算上轨和下轨
        upper_band = [m + num_std * s if not np.isnan(m) else np.nan for m, s in zip(middle_band, rolling_std)]
        lower_band = [m - num_std * s if not np.isnan(m) else np.nan for m, s in zip(middle_band, rolling_std)]
        
        return (middle_band, upper_band, lower_band)
    
    def calculate_all_indicators(self, price_data: Dict[str, Dict[str, float]], volume_data: Optional[Dict[str, int]] = None) -> Dict[str, Dict[str, float]]:
        """
        计算所有技术指标
        
        Args:
            price_data: 价格数据字典，键为日期，值为包含开盘价、收盘价、最高价、最低价的字典
            volume_data: 交易量数据字典，键为日期，值为交易量
            
        Returns:
            包含所有技术指标的字典，键为日期，值为包含各指标值的字典
        """
        logger.info("计算所有技术指标")
        
        try:
            # 按日期排序
            sorted_dates = sorted(price_data.keys())
            
            # 提取收盘价序列
            close_prices = [price_data[date]["close"] for date in sorted_dates]
            
            # 计算各种技术指标
            ma5 = self.calculate_ma(close_prices, 5)
            ma10 = self.calculate_ma(close_prices, 10)
            ma20 = self.calculate_ma(close_prices, 20)
            ma60 = self.calculate_ma(close_prices, 60)
            
            rsi = self.calculate_rsi(close_prices)
            
            macd_line, signal_line, histogram = self.calculate_macd(close_prices)
            
            middle_band, upper_band, lower_band = self.calculate_bollinger_bands(close_prices)
            
            # 组织结果
            indicators = {}
            for i, date in enumerate(sorted_dates):
                indicators[date] = {
                    "ma5": ma5[i],
                    "ma10": ma10[i],
                    "ma20": ma20[i],
                    "ma60": ma60[i],
                    "rsi": rsi[i],
                    "macd": macd_line[i],
                    "macd_signal": signal_line[i],
                    "macd_histogram": histogram[i],
                    "bb_middle": middle_band[i],
                    "bb_upper": upper_band[i],
                    "bb_lower": lower_band[i]
                }
            
            return indicators
        
        except Exception as e:
            logger.error(f"计算技术指标时出错: {str(e)}")
            return {}


class NewsSentimentAnalyzer:
    """新闻情感分析器"""
    
    def __init__(self, model_path: str):
        """
        初始化新闻情感分析器
        
        Args:
            model_path: 情感分析模型的路径
        """
        self.model_path = model_path
        logger.info(f"初始化新闻情感分析器，模型路径: {model_path}")
        # TODO: 加载实际的情感分析模型
    
    def analyze_batch(self, news_items: List[Dict[str, Any]]) -> Dict[str, Dict[str, float]]:
        """
        批量分析新闻情感
        
        Args:
            news_items: 新闻项列表，每个项是一个字典，包含'date'、'title'、'content'等
            
        Returns:
            情感分析结果字典，键为日期，值为包含情感得分和置信度的字典
        """
        logger.info(f"批量分析 {len(news_items)} 条新闻")
        
        try:
            # TODO: 实现实际的情感分析逻辑，这里使用模拟数据
            sentiment_results = {}
            
            for news in news_items:
                date = news.get("date", "")
                if not date:
                    continue
                
                # 模拟情感分析结果
                sentiment_score = np.random.uniform(-1.0, 1.0)
                confidence = np.random.uniform(0.6, 0.95)
                
                # 如果这一天已经有结果，取平均值
                if date in sentiment_results:
                    existing_score = sentiment_results[date]["sentiment_score"]
                    existing_confidence = sentiment_results[date]["confidence"]
                    new_score = (existing_score + sentiment_score) / 2
                    new_confidence = (existing_confidence + confidence) / 2
                    sentiment_results[date] = {
                        "sentiment_score": new_score,
                        "confidence": new_confidence
                    }
                else:
                    sentiment_results[date] = {
                        "sentiment_score": sentiment_score,
                        "confidence": confidence
                    }
            
            return sentiment_results
        
        except Exception as e:
            logger.error(f"分析新闻情感时出错: {str(e)}")
            return {}


class DatabaseConnector:
    """数据库连接器"""
    
    def __init__(self, connection_string: str):
        """
        初始化数据库连接器
        
        Args:
            connection_string: 数据库连接字符串
        """
        self.connection_string = connection_string
        logger.info("初始化数据库连接器")
        # TODO: 实际建立数据库连接
    
    def get_stock_news(self, stock_code: str, start_date: str, end_date: str) -> List[Dict[str, Any]]:
        """
        获取股票相关新闻
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            新闻项列表，每个项是一个字典，包含'date'、'title'、'content'等
        """
        logger.info(f"获取股票 {stock_code} 从 {start_date} 到 {end_date} 的相关新闻")
        
        # TODO: 实现实际的数据库查询逻辑，这里使用模拟数据
        try:
            # 模拟每天有0-3条新闻
            date_range = pd.date_range(start=start_date, end=end_date)
            
            news_items = []
            for date in date_range:
                date_str = date.strftime("%Y-%m-%d")
                num_news = np.random.randint(0, 4)  # 0到3条新闻
                
                for i in range(num_news):
                    # 模拟新闻数据
                    news_items.append({
                        "date": date_str,
                        "title": f"模拟新闻标题 {i+1} 关于股票 {stock_code}",
                        "content": f"这是一条关于股票 {stock_code} 的模拟新闻内容，第 {i+1} 条，日期为 {date_str}。",
                        "source": "模拟新闻源",
                        "url": f"https://example.com/news/{stock_code}/{date_str}/{i+1}"
                    })
            
            return news_items
        
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 相关新闻时出错: {str(e)}")
            return []


class DataIntegrator:
    """数据源整合器，负责集成各种数据源的数据"""
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化数据源整合器
        
        Args:
            config: 配置字典，包含API密钥、连接字符串等
        """
        self.price_data_source = PriceDataSource(config.get('price_api_key', ''))
        self.volume_data_source = VolumeDataSource(config.get('volume_api_key', ''))
        self.indicators_calculator = TechnicalIndicatorsCalculator()
        self.sentiment_analyzer = NewsSentimentAnalyzer(config.get('sentiment_model_path', ''))
        self.db_connector = DatabaseConnector(config.get('db_connection_string', ''))
        
        logger.info("初始化数据源整合器")
    
    def collect_historical_price_data(self, stock_codes: List[str], start_date: str, end_date: str) -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        收集多只股票的历史价格数据
        
        Args:
            stock_codes: 股票代码列表
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            价格数据字典，键为股票代码，值为日期到价格数据的映射
        """
        logger.info(f"收集 {len(stock_codes)} 只股票的历史价格数据")
        
        price_data = {}
        for code in stock_codes:
            price_data[code] = self.price_data_source.get_daily_prices(
                code, start_date, end_date
            )
        
        return price_data
    
    def collect_trading_volume_data(self, stock_codes: List[str], start_date: str, end_date: str) -> Dict[str, Dict[str, int]]:
        """
        收集多只股票的交易量数据
        
        Args:
            stock_codes: 股票代码列表
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            交易量数据字典，键为股票代码，值为日期到交易量的映射
        """
        logger.info(f"收集 {len(stock_codes)} 只股票的交易量数据")
        
        volume_data = {}
        for code in stock_codes:
            volume_data[code] = self.volume_data_source.get_daily_volumes(
                code, start_date, end_date
            )
        
        return volume_data
    
    def calculate_technical_indicators(self, price_data: Dict[str, Dict[str, Dict[str, float]]], volume_data: Dict[str, Dict[str, int]]) -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        计算技术指标
        
        Args:
            price_data: 价格数据字典，键为股票代码，值为日期到价格数据的映射
            volume_data: 交易量数据字典，键为股票代码，值为日期到交易量的映射
            
        Returns:
            技术指标字典，键为股票代码，值为日期到指标数据的映射
        """
        logger.info("计算技术指标")
        
        indicators = {}
        for code in price_data:
            indicators[code] = self.indicators_calculator.calculate_all_indicators(
                price_data[code], volume_data.get(code, None)
            )
        
        return indicators
    
    def collect_news_sentiment(self, stock_codes: List[str], start_date: str, end_date: str) -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        收集新闻情感分析结果
        
        Args:
            stock_codes: 股票代码列表
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            情感分析结果字典，键为股票代码，值为日期到情感数据的映射
        """
        logger.info(f"收集 {len(stock_codes)} 只股票的新闻情感分析结果")
        
        sentiment_data = {}
        for code in stock_codes:
            news_items = self.db_connector.get_stock_news(code, start_date, end_date)
            sentiment_data[code] = self.sentiment_analyzer.analyze_batch(news_items)
        
        return sentiment_data
    
    def integrate_all_data(self, stock_codes: List[str], start_date: str, end_date: str) -> Dict[str, Dict[str, Dict[str, Any]]]:
        """
        整合所有数据源
        
        Args:
            stock_codes: 股票代码列表
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            整合后的数据字典，键为股票代码，值为日期到所有特征数据的映射
        """
        logger.info(f"整合所有数据源: {len(stock_codes)} 只股票，从 {start_date} 到 {end_date}")
        
        # 收集各类数据
        price_data = self.collect_historical_price_data(stock_codes, start_date, end_date)
        volume_data = self.collect_trading_volume_data(stock_codes, start_date, end_date)
        indicators = self.calculate_technical_indicators(price_data, volume_data)
        sentiment_data = self.collect_news_sentiment(stock_codes, start_date, end_date)
        
        # 按日期整合
        integrated_data = {}
        for code in stock_codes:
            integrated_data[code] = self._align_time_series(
                price_data.get(code, {}),
                volume_data.get(code, {}),
                indicators.get(code, {}),
                sentiment_data.get(code, {})
            )
        
        return integrated_data
    
    def _align_time_series(self, price_data: Dict[str, Dict[str, float]], volume_data: Dict[str, int], 
                          indicators: Dict[str, Dict[str, float]], sentiment_data: Dict[str, Dict[str, float]]) -> Dict[str, Dict[str, Any]]:
        """
        对齐不同时间序列的数据
        
        Args:
            price_data: 价格数据，键为日期
            volume_data: 交易量数据，键为日期
            indicators: 技术指标数据，键为日期
            sentiment_data: 情感分析数据，键为日期
            
        Returns:
            对齐后的时间序列数据，键为日期，值为所有特征的字典
        """
        logger.info("对齐不同时间序列的数据")
        
        # 获取所有日期
        all_dates = set()
        all_dates.update(price_data.keys())
        all_dates.update(volume_data.keys())
        all_dates.update(indicators.keys())
        all_dates.update(sentiment_data.keys())
        
        # 对齐数据
        aligned_data = {}
        for date in sorted(all_dates):
            # 只有当价格数据可用时才包含这一天
            if date in price_data:
                # 初始化该日期的数据
                daily_data = {}
                
                # 添加价格数据
                daily_data.update(price_data[date])
                
                # 添加交易量数据
                if date in volume_data:
                    daily_data["volume"] = volume_data[date]
                
                # 添加技术指标数据
                if date in indicators:
                    for key, value in indicators[date].items():
                        daily_data[key] = value
                
                # 添加情感分析数据
                if date in sentiment_data:
                    for key, value in sentiment_data[date].items():
                        daily_data[f"sentiment_{key}"] = value
                
                # 保存对齐后的数据
                aligned_data[date] = daily_data
        
        return aligned_data 